import torch
import configs
from models import *
from dataset import *
import torch.nn.functional as F

first_run = False
loss_simple_method = True

if __name__ == '__main__':
    config = configs.Config()
    config.device = 'cuda' if torch.cuda.is_available() else 'cpu'

    if first_run:
        filter_useful_data(config)
        Word2Vec_transformer(config)
        Gensim_model = Embedding_pretrained(config)
        user_count, item_count, vocab_num = Ordered_transformer(config)
        config.num_user = user_count
        config.num_item = item_count
        config.vocab_num = vocab_num

    print('Config argument: \n', config)

    Amazon_loader = Amazon_DataLoader(config)

    # train PMF model
    PMF_model = PMF(config)
    PMF_model.weight_to(config.device)
    for epoch in range(config.PMF_epoch):
        print('Train PMF epoch [{}/{}]'.format(epoch, config.PMF_epoch))
        for i, input_dict in enumerate(Amazon_loader):
            error = 0.0
            for sample in range(input_dict['user_id'].shape[0]):
                error += PMF_model.train_step(
                    input_dict['user_id'][sample],
                    input_dict['item_id'][sample],
                    input_dict['rating_num'][sample],
                )
            error /= input_dict['user_id'].shape[0]
            if i % config.show_step == 0:
                print('iteration {}, error {}'.format(i, error))
    torch.save(PMF_model.state_dict(), config.PMF_weight_path)
    print('PMF model train over, save over!')

    # LSTM language model.
    LSTM_model = LSTM_pretrained(config).to(config.device)
    optimizer = torch.optim.Adam(LSTM_model.parameters(), lr=config.LSTM_LR)
    for epoch in range(config.LSTM_epoch):
        print('Train LSTM epoch [{}/{}]'.format(epoch, config.LSTM_epoch))
        for i, input_dict in enumerate(Amazon_loader):
            LSTM_model.update(input_dict)
            outputs = LSTM_model()
            loss = LSTM_model.loss_fn()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if i % config.show_step == 0:
                print('iteration {}, loss {}'.format(i, loss.item()))
    torch.save(LSTM_model.state_dict(), config.LSTM_weight_path)
    print('LSTM model train over, save over!')

    # Dual model train.
    Review2Rating_model = Review2Rating(config).to(config.device)
    optimizer_review2rating = torch.optim.SGD(
        Review2Rating_model.parameters(), lr=config.Review2Rating_LR
    )

    Rating2Review_model = Rating2Review(config).to(config.device)
    optimizer_rating2review = torch.optim.SGD(
        Rating2Review_model.parameters(), lr=config.Rating2Review_LR
    )

    # The embedding is still fake!
    # dual loss has bug to -inf!

    for epoch in range(config.dual_model_epoch):
        print('Train dual model epoch [{}/{}]'.format(epoch, config.dual_model_epoch))
        for i, input_dict in enumerate(Amazon_loader):
            for key in input_dict:
                if input_dict[key] is None:
                    continue
                input_dict[key] = input_dict[key].to(config.device)
                input_dict[key].requires_grad = False

            with torch.no_grad():
                # to make the extend data.
                PMF_predict = PMF_model.batch_forward(input_dict['user_id'], input_dict['item_id'])
                PMF_predict = PMF_predict.detach()
                # print('PMF predict device ', PMF_predict.device)  # for debug.

                LSTM_model.update(input_dict)
                LSTM_predict = LSTM_model()
                sentence_embedding = LSTM_model.get_sentence_embedding().detach()
                input_dict['sentence_embedding'] = sentence_embedding.squeeze(0).to(config.device)

            # To optimize Review to Rating model
            Review2Rating_model.update(input_dict)
            rating_predict = Review2Rating_model()
            l1_loss = Review2Rating_model.loss_function()

            Rating2Review_model.update(input_dict)
            review_predict = Rating2Review_model()
            l2_loss = Rating2Review_model.loss_function()

            # print('iteration {}: '.format(i),
            #       l2_loss,
            #       'Rating min max: ',
            #       input_dict['rating_num'].min(),
            #       input_dict['rating_num'].max(),
            #       input_dict['rating_num'],
            #       'Embedding weight: ',
            #       Rating2Review_model.rating_embedding.weight.min(),
            #       Rating2Review_model.rating_embedding.weight.max(),
            #       'Sentence index: ',
            #       )
            # if torch.isnan(Rating2Review_model.rating_embedding.weight).any():
            #     review_predict = Rating2Review_model(debug=True)
            #     # print(input_dict)
            #     for key in input_dict:
            #         print('key {}, min {}, max {}'.format(key, input_dict[key].min(), input_dict[key].max()))
            #     print(review_predict)
            #     print(review_predict.min(), review_predict.max())
            #
            #     if torch.isnan(l2_loss):
            #         exit(0)

            dual_loss = \
                - torch.sum((input_dict['rating_num'].to(config.device)-PMF_predict)**2 / (2 * config.vars)) \
                + F.cross_entropy(
                    review_predict.view(-1, config.vocab_num),
                    input_dict['target_index'].view(-1).to(config.device),
                ) \
                - F.cross_entropy(
                    LSTM_predict.view(-1, config.vocab_num),
                    input_dict['target_index'].view(-1).to(config.device),
                ) \
                + torch.sum((input_dict['rating_num'].to(config.device)-rating_predict)**2 / (2 * config.vars))

            # One possible method to simplify the backward is to calculate the sum of both loss.
            if loss_simple_method:
                optimizer_review2rating.zero_grad()
                optimizer_rating2review.zero_grad()
                loss_all = l1_loss + l2_loss + (config.lambda_review2rating + config.lambda_rating2review) * dual_loss
                loss_all.backward()
                # print('loss l1 {}, l2 {}, dual {}'.format(l1_loss, l2_loss, dual_loss))
                # print('grad of fuse layer: ', Rating2Review_model.rating_embedding.weight.grad)
                optimizer_review2rating.step()
                optimizer_rating2review.step()
            else:
                # This way still have bug!
                optimizer_review2rating.zero_grad()
                loss_review2rating = l1_loss + config.lambda_review2rating * dual_loss
                loss_review2rating.backward(retain_graph=True)  # retain graph is to save the calc graph for second use.
                optimizer_review2rating.step()

                Review2Rating_model.update(input_dict)
                rating_predict = Review2Rating_model()
                l1_loss = Review2Rating_model.loss_function()

                Rating2Review_model.update(input_dict)
                review_predict = Rating2Review_model()
                l2_loss = Rating2Review_model.loss_function()

                dual_loss = \
                    - torch.sum((input_dict['rating_num'].to(config.device) - PMF_predict) ** 2 / (2 * config.vars)) \
                    + F.cross_entropy(
                        review_predict.view(-1, config.vocab_num),
                        input_dict['target_index'].view(-1).to(config.device),
                    ) \
                    - F.cross_entropy(
                        LSTM_predict.view(-1, config.vocab_num),
                        input_dict['target_index'].view(-1).to(config.device),
                    ) \
                    + torch.sum((input_dict['rating_num'].to(config.device) - rating_predict) ** 2 / (2 * config.vars))

                optimizer_rating2review.zero_grad()
                loss_rating2review = l2_loss + config.lambda_rating2review * dual_loss
                loss_rating2review.backward()  # after this usage, the graph can be delete.
                optimizer_rating2review.step()

            # To watch the train iteration
            if i % config.show_step == 0:
                print('iteration {}, l1 loss {}, l2 loss {}, dual loss {}'.format(
                    i, l1_loss.item(), l2_loss.item(), dual_loss.item())
                )

    torch.save(Review2Rating_model.state_dict(), config.Review2Rating_weight_path)
    torch.save(Rating2Review_model.state_dict(), config.Rating2Review_weight_path)
    print('Dual model train over, save over!')

    print('Train over!')
